package main

import (
	"text/template"
	"unicode"

	"entgo.io/ent/entc"
	"entgo.io/ent/entc/gen"
)

func main() {
	// 创建自定义函数映射
	funcs := template.FuncMap{
		"toSnakeCase": ToSnakeCase, // 注册蛇形命名转换函数
		"export": func(s string) string { // 将字段名首字母大写用于 struct/Set 方法
			if s == "" {
				return ""
			}
			r := []rune(s)
			r[0] = unicode.ToUpper(r[0])
			return string(r)
		},
	}
	// 将模板放到变量中，供 config 和扩展共享
	templates := []*gen.Template{
		gen.MustParse(gen.NewTemplate("static"), nil),
		gen.MustParse(
			gen.NewTemplate("repository").
				Funcs(funcs).                             // 注入自定义函数
				ParseFiles("templates/repository.tmpl")), // 解析你的模板文件
	}

	cfg := &gen.Config{
		Package: "entdemo/ent",
		Target:  "ent",
		Schema:  "schema",
		Features: []gen.Feature{
			gen.FeatureUpsert,
			gen.FeatureVersionedMigration,
			gen.FeatureModifier,
		},
	}

	// 将自定义扩展通过 entc.Extensions 注入到 Generate，指定自定义输出目录并传入模板
	var repoTmpl *gen.Template
	if len(templates) > 1 {
		repoTmpl = templates[1]
	}
	err := entc.Generate("./schema", cfg, entc.Extensions(NewRepoExtension("entdemo", "pkg/repository", repoTmpl)))
	if err != nil {
		panic(err)
	}
}

// ToSnakeCase 将驼峰命名转换为蛇形命名（例如 "UserName" -> "user_name"）
func ToSnakeCase(s string) string {
	var result []rune
	for i, r := range s {
		if unicode.IsUpper(r) {
			if i > 0 {
				result = append(result, '_')
			}
			result = append(result, unicode.ToLower(r))
		} else {
			result = append(result, r)
		}
	}
	return string(result)
}
